package com.zengde.ai_sandbox.service.impl;

import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionRequest;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessage;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessageRole;
import com.volcengine.ark.runtime.service.ArkService;
import com.zengde.ai_sandbox.service.SuperBigModelInterface;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.time.Duration;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Service("doubaoService")
@Slf4j
public class DoubaoBigModelServiceImpl implements SuperBigModelInterface {
    
    private final ArkService arkService;
    private final Map<String, List<ChatMessage>> sessionMessages;
    private final boolean useMockService;
    
    @Value("${my.doubao.normal.model.name}")
    private String normalModelName;
    
    @Value("${my.doubao.think.model.name}")
    private String thinkModelName;
    
    @Value("${my.doubao.image.model.name}")
    private String imageModelName;
    
    public DoubaoBigModelServiceImpl() {
        String apiKey = System.getenv("ARK_API_KEY");
        
        // 如果没有配置API密钥，则使用模拟服务
        if (apiKey == null || apiKey.isEmpty()) {
            log.warn("未配置ARK_API_KEY环境变量，将使用模拟服务");
            this.arkService = null;
            this.useMockService = true;
        } else {
            this.arkService = ArkService.builder().apiKey(apiKey).build();
            this.useMockService = false;
        }
        
        this.sessionMessages = new ConcurrentHashMap<>();
    }
    
    /**
     * 通用方法：非流式调用
     */
    private String chatCommon(String sessionId, String content, String modelName) {
        if (useMockService) {
            return mockChatResponse(sessionId, content, modelName);
        }
        
        List<ChatMessage> messages = getOrCreateSession(sessionId);
        messages.add(ChatMessage.builder()
                .role(ChatMessageRole.USER)
                .content(content)
                .build());

        ChatCompletionRequest request = ChatCompletionRequest.builder()
                .model(modelName)
                .messages(messages)
                .build();

        StringBuilder responseBuilder = new StringBuilder();
        arkService.streamChatCompletion(request)
                .doOnNext(choice -> {
                    if (choice.getChoices().size() > 0) {
                        responseBuilder.append(String.valueOf(choice.getChoices().get(0).getMessage().getContent()));
                    }
                })
                .blockingLast();  // 等待流完成

        String response = responseBuilder.toString();
        
        // 保存助手的回复到会话历史
        messages.add(ChatMessage.builder()
                .role(ChatMessageRole.ASSISTANT)
                .content(response)
                .build());

        return response;
    }
    
    /**
     * 通用方法：流式调用
     */
    private Flux<String> streamChatCommon(String sessionId, String content, String modelName) {
        if (useMockService) {
            return mockStreamChatResponse(sessionId, content, modelName);
        }
        
        List<ChatMessage> messages = getOrCreateSession(sessionId);
        messages.add(ChatMessage.builder()
                .role(ChatMessageRole.USER)
                .content(content)
                .build());

        ChatCompletionRequest request = ChatCompletionRequest.builder()
                .model(modelName)
                .messages(messages)
                .build();

        StringBuilder fullResponse = new StringBuilder();
        
        return Flux.from(arkService.streamChatCompletion(request))
                .map(choice -> {
                    if (choice.getChoices().size() > 0) {
                        String messageContent = String.valueOf(choice.getChoices().get(0).getMessage().getContent());
                        fullResponse.append(messageContent);
                        return messageContent;
                    }
                    return "";
                })
                .doOnComplete(() -> {
                    // 保存完整的助手回复到会话历史
                    messages.add(ChatMessage.builder()
                            .role(ChatMessageRole.ASSISTANT)
                            .content(fullResponse.toString())
                            .build());
                });
    }
    
    /**
     * 模拟聊天响应
     */
    private String mockChatResponse(String sessionId, String content, String modelName) {
        log.info("使用模拟服务: sessionId={}, content={}, modelName={}", sessionId, content, modelName);
        
        List<ChatMessage> messages = getOrCreateSession(sessionId);
        messages.add(ChatMessage.builder()
                .role(ChatMessageRole.USER)
                .content(content)
                .build());
        
        // 根据不同的模型生成不同的响应
        String response;
        if (modelName.contains("thinking")) {
            response = "这是一个模拟的深度思考回复。我会先分析一下您的问题，然后给出一个详细的回答。\n\n" +
                    "您说的是：" + content + "\n\n" +
                    "基于您的问题，我认为...(思考中)...最合适的解决方案是...";
        } else if (modelName.contains("vision")) {
            response = "这是一个模拟的图像分析回复。我注意到您上传的图片中包含...(分析中)...";
        } else {
            response = "您好，这是一个模拟的AI回复。您问的是：" + content + "\n\n" +
                    "由于当前使用的是模拟服务，无法提供真实的AI回答。请配置ARK_API_KEY环境变量后再试。";
        }
        
        // 保存助手的回复到会话历史
        messages.add(ChatMessage.builder()
                .role(ChatMessageRole.ASSISTANT)
                .content(response)
                .build());
        
        return response;
    }
    
    /**
     * 模拟流式聊天响应
     */
    private Flux<String> mockStreamChatResponse(String sessionId, String content, String modelName) {
        log.info("使用模拟流式服务: sessionId={}, content={}, modelName={}", sessionId, content, modelName);
        
        List<ChatMessage> messages = getOrCreateSession(sessionId);
        messages.add(ChatMessage.builder()
                .role(ChatMessageRole.USER)
                .content(content)
                .build());
        
        // 根据不同的模型生成不同的响应
        String[] chunks;
        if (modelName.contains("thinking")) {
            chunks = new String[] {
                    "这是一个模拟的",
                    "深度思考",
                    "回复。",
                    "我会先分析一下您的问题，",
                    "然后给出一个详细的回答。\n\n",
                    "您说的是：" + content + "\n\n",
                    "基于您的问题，",
                    "我认为...",
                    "(思考中)...",
                    "最合适的解决方案是..."
            };
        } else {
            chunks = new String[] {
                    "您好，",
                    "这是一个",
                    "模拟的AI回复。",
                    "您问的是：" + content + "\n\n",
                    "由于当前使用的是模拟服务，",
                    "无法提供真实的AI回答。",
                    "请配置ARK_API_KEY环境变量后再试。"
            };
        }
        
        // 构建完整响应
        StringBuilder fullResponse = new StringBuilder();
        for (String chunk : chunks) {
            fullResponse.append(chunk);
        }
        
        final String finalResponse = fullResponse.toString();
        
        // 模拟流式响应
        return Flux.fromArray(chunks)
                .delayElements(Duration.ofMillis(200))
                .doOnComplete(() -> {
                    // 保存完整的助手回复到会话历史
                    messages.add(ChatMessage.builder()
                            .role(ChatMessageRole.ASSISTANT)
                            .content(finalResponse)
                            .build());
                });
    }

    @Override
    public String ntns(String sessionId, String content) {
        return chatCommon(sessionId, content, normalModelName);
    }

    @Override
    public Flux<String> nts(String sessionId, String content) {
        return streamChatCommon(sessionId, content, normalModelName);
    }

    @Override
    public Flux<String> ts(String sessionId, String content) {
        return streamChatCommon(sessionId, content, thinkModelName);
    }

    @Override
    public String tns(String sessionId, String content) {
        return chatCommon(sessionId, content, thinkModelName);
    }

    @Override
    public String imgChat(String sessionId, String content, String imageUrl) {
        if (useMockService) {
            return "这是一个模拟的图像聊天回复。我看到了您上传的图片：" + imageUrl + "，以及您的问题：" + content;
        }
        
        List<ChatMessage> messages = getOrCreateSession(sessionId);
        
        // 构建包含图像的消息
        // 注意：根据volcengine SDK 0.2.13版本，可能不支持contentList方法
        // 这里改用content方法，将图像URL和文本内容合并
        String combinedContent = String.format("图片URL: %s\n\n%s", imageUrl, content);
        
        messages.add(ChatMessage.builder()
                .role(ChatMessageRole.USER)
                .content(combinedContent)
                .build());

        ChatCompletionRequest request = ChatCompletionRequest.builder()
                .model(imageModelName)
                .messages(messages)
                .build();

        StringBuilder responseBuilder = new StringBuilder();
        arkService.streamChatCompletion(request)
                .doOnNext(choice -> {
                    if (choice.getChoices().size() > 0) {
                        responseBuilder.append(String.valueOf(choice.getChoices().get(0).getMessage().getContent()));
                    }
                })
                .blockingLast();  // 等待流完成

        String response = responseBuilder.toString();
        
        // 保存助手的回复到会话历史
        messages.add(ChatMessage.builder()
                .role(ChatMessageRole.ASSISTANT)
                .content(response)
                .build());

        return response;
    }

    @Override
    public void initializeSession(String sessionId, String rolePrompt) {
        log.info("初始化会话: sessionId={}, rolePrompt={}", sessionId, rolePrompt);
        List<ChatMessage> messages = new ArrayList<>();
        messages.add(ChatMessage.builder()
                .role(ChatMessageRole.SYSTEM)
                .content(rolePrompt)
                .build());
        sessionMessages.put(sessionId, messages);
    }

    private List<ChatMessage> getOrCreateSession(String sessionId) {
        return sessionMessages.computeIfAbsent(sessionId, k -> new ArrayList<>());
    }
}
